import os.path
from datetime import timedelta

import pandas as pd
from tqdm import tqdm
from xbase_util.common_util import build_es_expression, s2date

from src.bean.splunk_config_enum import SplunkConfigEnum
from src.col_bean.dangerous_enum import ColDangerousEnum
from src.constant import project_root_path, expression, expression_tag
from src.extract.session import get_session
from src.util.common_util import printx, get_all_threats
from src.util.config_manager import ConfigManager


def extract_normal(config: ConfigManager, start_time, end_time, exp, default_save_file):
    print(f"开始获取正常数据:{start_time}-{end_time}")
    exp = build_es_expression(size=5000, start_time=s2date(start_time), end_time=s2date(end_time),
                              arkime_expression=exp)
    dangerous_all = get_all_threats(host=SplunkConfigEnum.host.value, port=SplunkConfigEnum.port.value,
                                    username=SplunkConfigEnum.username.value,
                                    password=SplunkConfigEnum.password.value, scheme=SplunkConfigEnum.scheme.value,
                                    start_time=s2date(start_time), end_time=s2date(end_time), use_df=False)
    duration = 0
    count = 0
    session_list = get_session(exp, config)
    if session_list is None or len(session_list) == 0:
        print("session_list is empty")
        return
    with tqdm(total=len(dangerous_all)) as pbs:
        printx("session获取完毕，开始打标签")
        session_list['firstPacket_datetime'] = pd.to_datetime(session_list['firstPacket'], unit='ms')
        session_list['lastPacket_datetime'] = pd.to_datetime(session_list['firstPacket'], unit='ms')
        session_list['source.port'].astype(int)
        session_list['destination.port'].astype(int)
        for ind, d_item in enumerate(dangerous_all):
            if d_item['type'] == 'ids':
                sip = session_list['source.ip']
            else:
                sip = session_list['http.xffIp'].apply(
                    lambda x: x[0] if isinstance(x, list) and len(x) > 0 else "")
            threat_time = s2date(d_item['THREAT_TIME'])
            result_session = session_list[(
                    (sip == d_item["SIP"]) &
                    (session_list['source.port'] == int(d_item["S_PORT"])) &
                    (session_list['destination.ip'] == d_item["DIP"]) &
                    (session_list['destination.port'] == int(d_item["D_PORT"])) &
                    (session_list['firstPacket_datetime'] >= threat_time - timedelta(minutes=duration)) &
                    (session_list['lastPacket_datetime'] <= threat_time + timedelta(minutes=duration))
            )]
            count = count + len(result_session)
            if len(result_session) != 0:
                printx(
                    f'[打标签]匹配到：{len(result_session)},共匹配到：{count}')
                session_list.loc[session_list['id'].isin(result_session['id']), [
                    ColDangerousEnum.deny_method.value, ColDangerousEnum.protocol.value,
                    ColDangerousEnum.severity.value,
                    ColDangerousEnum.threat_summary.value, ColDangerousEnum.traffic_type.value,
                    ColDangerousEnum.is_dangerous.value]] = [
                    d_item['DENY_METHOD'],
                    d_item['PROTOCOL'],
                    d_item['SEVERITY'],
                    d_item['type'],
                    d_item['THREAT_SUMMARY'],
                    True
                ]
            pbs.set_description(f"[打标签]进度:{ind + 1:02d}/{len(dangerous_all)}")
            pbs.update()
    session_list.drop(columns=['firstPacket_datetime', 'lastPacket_datetime'], inplace=True)
    session_list = session_list[session_list[ColDangerousEnum.is_dangerous.value] != True]
    session_list.to_csv(os.path.join(project_root_path, 'data', default_save_file), index=False)


def extract_normal_entrance(config: ConfigManager) -> None:
    generate_am_pm_data([
        # '1', '2', '3', '4', '5', '6', '7', '8', '9', '10', '11', '12', '13', '14', '15', '16', '17'
        # ,'18','19','20','21','22','23','24','25','26','27','28'
    ], config)


def generate_am_pm_data(days, config):
    # year = 2025
    # month = 8
    # for day in days:
    #     for tag in ['am', 'pm', 'night', 'midnight']:
    #         if tag == 'am':
    #             start_time = f"{year}/{month:02d}/{day} 11:00:00"
    #             end_time = f"{year}/{month:02d}/{day} 11:40:00"
    #         elif tag == 'pm':
    #             start_time = f"{year}/{month:02d}/{day} 14:00:00"
    #             end_time = f"{year}/{month:02d}/{day} 17:30:00"
    #         elif tag == 'night':
    #             start_time = f"{year}/{month:02d}/{day} 17:30:00"
    #             end_time = f"{year}/{month:02d}/{day} 23:59:59"
    #         else:
    #             start_time = f"{year}/{month:02d}/{day} 00:00:00"
    #             end_time = f"{year}/{month:02d}/{day} 09:00:00"
    #         t = s2date(start_time)
    #         dir = os.path.join(project_root_path, 'data',
    #                            f"normal_{expression_tag}_{tag}_{t.month:02d}_{t.day:02d}.csv")
    #         if not os.path.exists(dir):
    #             extract_normal(config, start_time, end_time, expression[expression_tag], dir)
    #         else:
    #             printx(f"{dir}该文件存在")
    dir = os.path.join(project_root_path, 'data', f"scan.csv")
    extract_normal(config, "2025/08/18 08:53:02", "2025/08/18 09:53:02", expression[expression_tag],
                   dir)
